package com.atguigu.day08;

import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;

public class Example2 {
    public static void main(String[] args) {
        var env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        env.enableCheckpointing(10 * 1000L);

        env
                .addSource(new CounterSource())
                .print();

        try {
            env.execute();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static class CounterSource extends RichParallelSourceFunction<Long> implements CheckpointedFunction {
        private boolean running = true;
        private long offset = 0L;
        @Override
        public void run(SourceContext<Long> ctx) throws Exception {
            // 获取一把锁
            final Object lock = ctx.getCheckpointLock();

            while (running) {
                synchronized (lock) {
                    ctx.collect(offset);
                    offset += 1L;
                }

                Thread.sleep(1000L);
            }
        }

        @Override
        public void cancel() {
            running = false;
        }

        // 声明一个算子状态用来保存刚消费完的偏移量
        // 算子状态中没有ValueState
        // 所以只能使用列表状态变量保存一个值（刚消费完的偏移量）
        private ListState<Long> state;

        // 程序第一次启动时触发调用
        // 故障恢复时触发调用
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            state = context.getOperatorStateStore().getListState(
                    new ListStateDescriptor<>(
                            "offset",
                            Types.LONG
                    )
            );

            // 列表状态变量中要么只有一个元素（故障恢复），要么没有元素（第一次启动的时候）
            for (var l : state.get()) {
                offset = l;
            }
        }

        // 当source的并行子任务接收到 检查点分界线 时，触发调用
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            state.clear();
            state.add(offset);
        }
    }
}
